# 展示模型动画并保存视频
from os.path import dirname, join, abspath
import yaml
def vizPlay(model, viz, qs, dt, ROBOT = False):

    # 根据Robot关键字的bool值，选取特定的init.yml文件
    if ROBOT:
        # 这里open的基路径是main.py
        with open(join(dirname(str(abspath(__file__))),'robot_init.yml'), 'r', encoding='utf-8') as f:
            init_config= yaml.load(f.read(), Loader=yaml.FullLoader)
    else:
        with open(join(dirname(str(abspath(__file__))),'init.yml'), 'r', encoding='utf-8') as f:
            init_config= yaml.load(f.read(), Loader=yaml.FullLoader)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


    joint_names = init_config.get('joint_names',[])

    frame_ids = []
    for joint_name in joint_names:
        frame_ids.append(model.getFrameId(joint_name))

    def my_callback(*args):
        for frame_id in frame_ids:
            viz.drawFrameVelocities(frame_id)

    with viz.create_video_ctx(init_config['video_path']):
        viz.play(qs, dt, callback=my_callback)